# coding=utf-8
"""
 _moc_ast.yaml 

@license: License: LGPL
@copyright: (c) 2009, dogsing.cn
@author: ant-man(Bin Zhao)
"""

from cn.pigersing.moc.lang.mocConstant import mocInt, mocFloat, mocString, mocGNone, mocTrue, mocFalse
from cn.pigersing.moc.lang.mocInstance import newtemp
from cn.pigersing.moc.utils.langUtils import StaticNum
import cn.pigersing.moc.lang.mocCst as mocCst
import copy
import mocAst
import mocLexer as KW


# for visit_ID to identify the keywords
_keywords = [ v.lower() for v in KW.mocLexer.keywords]

class MocTranslator(mocAst.NodeVisitor):


    def __init__(self, out=None):
        """
        @out : output stream where to output intermediate code, which must contain the "write" method
        """
        self.out = out
        self.codelist = []
        self.labelno = StaticNum()

    def __reset(self):
        """
        reset current translator and save the state
        """
        ret = (self.codelist, self.labelno)
        self.codelist = []
        self.labelno = StaticNum()
        return ret

    def __restore(self, state):
        """
        restore the state
        """
        self.codelist, self.labelno = state

    def translate(self, ast):
        """
        Translate a AST subtree to intermediate code
        @param ast: subtree of a mocAST
        """
        if not ast : return []
        state = self.__reset()
        ast.accept(self)
        ret = self.codelist
        # restore previous state
        self.__restore(state)
        return ret

    def newlabel(self):
        return mocCst.Label(self.labelno(), 0)

    def translate_children(self, node):
        """ Performing transition the children """
        for n in node.children() :
            n.accept(self)

    def emitlabel(self, label):
        """ Generate a label for current operation
            @param label_no: label number
        """
        label.lineno = len(self.codelist)
        if self.out :
            self.out.write("%s" % label)

    def emit(self, code):
        """ Generate a new line of translated source
            @param code: code write to the default output device
        """
        self.codelist.append(code)
        if self.out :
            self.out.write("\t%s\n" % code)

    def emitjumps(self, test, t, f):
        """ Generate the jump instructions, according to the
            given parameters
            @param test: condition for 'if' instruction
            @param t: true point to jump, if 0 then fall
            @param f: false point to jump, if 0 then fall
        """
        if t != 0 and f != 0 :
            self.emit(mocCst.IfTrue(test, mocCst.Goto(t)))
            self.emit(mocCst.Goto(f))
        elif t != 0 :
            self.emit(mocCst.IfTrue(test, mocCst.Goto(t)))
        elif f != 0:
            self.emit(mocCst.IfFalse(test, mocCst.Goto(f)))

    def visit_Constant(self, node):
        if node.type == "int" :
            const = mocInt.new_inst(node.value)
        elif node.type == "float" :
            const = mocFloat.new_inst(node.value)
        elif node.type == "bool" :
            const = mocTrue if node.value else mocFalse
        elif node.type == "string" :
            const = mocString.new_inst(node.value)
        else :
            const = mocGNone
        return const

    def visit_ID(self, node):
        """@return: identifier name"""
        if node.name in _keywords :
            raise ValueError("Identifier conflicts with keyword %s %s" % (node.name, node.coord))
        return node.name

    def visit_UnaryOp(self, node):
        # first translate the expression on 
        # the right hand side of operand
        if node.op in ["+", "-", "~", "!"] :
            temp = newtemp()
            expr = node.expr.accept(self)
            self.emit(mocCst.UnaryOp(node.op, expr, temp))
            return temp
        elif node.op in ["p++", "p--"] :
            temp = newtemp()
            mocAst.Assignment("=", mocAst.ID(temp), node.expr).accept(self)
            rexpr = copy.copy(node.expr)
            mocAst.Assignment("=", node.expr, mocAst.ArithOp(node.op[-1], rexpr, mocAst.Constant("int", 1))).accept(self)
            return temp
        elif node.op in ["++", "--"] :
            rexpr = copy.copy(node.expr)
            mocAst.Assignment("=", node.expr, mocAst.ArithOp(node.op[-1], rexpr, mocAst.Constant("int", 1))).accept(self)
            return node.expr.accept(self)
        else :
            raise ValueError("This '%s' should be happen in visit_UnaryOp()" % node.op)

    def visit_ArithOp(self, node):
        lhs = node.left.accept(self)
        rhs = node.right.accept(self)
        # generate a temporary variable to hold the new value
        temp = newtemp()
        self.emit(mocCst.BinaryOp(node.op, lhs, rhs, temp))
        return temp

    def visit_RelOp(self, node):
        if hasattr(node, "trueL") and hasattr(node, "falseL") :
            lhs = node.left.accept(self)
            rhs = node.right.accept(self)
            # generate the following test condition
            # lhs op rhs, eg: a == b
            test = mocCst.BinaryOp(node.op, lhs, rhs, None)
            # the "trueL" and "falseL" is the label number, inherit from parent
            t = node.trueL
            f = node.falseL
            self.emitjumps(test, t, f)
        else :
            assignment = mocAst.Assignment("=", mocAst.ID(newtemp()), node)
            return assignment.accept(self)

    def visit_LogicalOp(self, node):
        if hasattr(node, "trueL") and hasattr(node, "falseL") :
            # the "trueL" and "falseL" is the label number, inherit from parent
            t = node.trueL
            f = node.falseL
            gen_label = False
            if node.op == "||" :
                node.left.trueL = self.newlabel() if t == 0 else t
                gen_label = t == 0
                node.left.falseL = 0
            elif node.op == "&&" :
                node.left.falseL = self.newlabel() if f == 0 else f
                gen_label = f == 0
                node.left.trueL = 0
            node.right.trueL = t
            node.right.falseL = f
            node.left.accept(self)
            node.right.accept(self)
            if gen_label :
                self.emitlabel(node.left.trueL if t == 0 else node.left.falseL)
        else :
            assignment = mocAst.Assignment("=", mocAst.ID(newtemp()), node)
            return assignment.accept(self)

    def visit_Assignment(self, node):
        # each type of assignment will always return the lhs
        lhs = None
        if isinstance(node.lvalue, mocAst.Constant) and node.op != ":":
            # if a constant value like a integer at left hand side of the assignment
            # this should not be allowed
            raise ValueError("Constant expression in left hands of assignment is illegal!")
        if node.op == ":" :
            # this token may be used in a declaration of a dictionary,
            # say: "dict = {key = val, key2 = val2} ;"
            # or in a slicing operation, say: "some = list[a:b] ;"
            # the action is decided by the visit_ParamList()
            if isinstance(node.lvalue, mocAst.Assignment) or isinstance(node.rvalue, mocAst.Assignment) :
                raise ValueError("slice can not be resolved % s" % node.coord)
            if isinstance(node.rvalue, mocAst.FuncDef) :
                rhs = mocAst.Assignment("=", mocAst.ID(newtemp()), node.rvalue).accept(self)
            else:
                rhs = node.rvalue.accept(self)
            lhs = node.lvalue.accept(self)
            if node.actionT == "list" :
                self.emit(mocCst.Range(lhs, rhs))
            else :
                self.emit(mocCst.Dictitem(lhs, rhs))
        elif isinstance(node.lvalue, mocAst.ObjectRef) :
            #  when a object is at left hands side, the assignment expression, may bring new field to a object
            temp = newtemp()
            node.lvalue.assign = True
            node.lvalue.temp = temp
            mocAst.Assignment("=", mocAst.ID(temp), node.rvalue, node.coord).accept(self)
            lhs = node.lvalue.accept(self)
        elif isinstance(node.lvalue, mocAst.CollectionRef) :
            #  when a object is at left hands side, the assignment expression, may bring new field to a object
            temp = newtemp()
            node.lvalue.assign = True
            mocAst.Assignment("=", mocAst.ID(temp), node.rvalue, node.coord).accept(self)
            self.emit(mocCst.Setret(temp))
            lhs = node.lvalue.accept(self)
        elif isinstance(node.rvalue, (mocAst.LogicalOp, mocAst.RelOp)) :
            # if the right hand value is a logical expression or a relation expression
            # then we must specify its true entry and false entry
            t = self.newlabel()
            f = self.newlabel()
            # set the inherit attributes to child node
            node.rvalue.trueL = t
            node.rvalue.falseL = f
            # transfer from children
            node.rvalue.accept(self)
            lhs = node.lvalue.accept(self)
            # left hand value now equals true
            self.emitlabel(t) ; self.emit(mocCst.AssignOp(mocTrue, lhs))
            # generate a goto instruction jumping the false assignment
            label = self.newlabel()
            self.emit(mocCst.Goto(label))
            # left hand value now equals false
            self.emitlabel(f) ; self.emit(mocCst.AssignOp(mocFalse, lhs))
            # generate a label so when the value equals true, program can jump
            # to the next line
            self.emitlabel(label)
        elif isinstance(node.rvalue, mocAst.FuncDef) :
            param_code, code_list = node.rvalue.accept(self)
            lhs = node.lvalue.accept(self)
            self.emit(mocCst.CallableDef(param_code, code_list, lhs))
        else :
            rhs = node.rvalue.accept(self)
            lhs = node.lvalue.accept(self)
            if node.op in ["+=", "-=", "*=", "/=", "%=", "<<=", ">>=", "~=", "&=", "|="] :
                # this just convert the "lhs @= rhs" to " lhs = lhs @ rhs" 
                self.emit(mocCst.BinaryOp(node.op[0], lhs, rhs, lhs))
            else :
                # the ordinary assignment operand
                self.emit(mocCst.AssignOp(rhs, lhs))
        return lhs

    def visit_ParamList(self, node):
        if isinstance(node.parent, mocAst.FuncDef):
            # parsing FuncDef, translate each parameters in (name, code) 
            hasdefault = False # has default parameters
            for param in node.params :
                if isinstance(param, mocAst.ID):
                    # this means we have encountered a default parameter then follow another normal
                    if hasdefault : raise ValueError("Default parameters must appear in the end of parameter list")
                    # Translate and put it into FunDef's param_code
                    name = param.name
                    code_list = None
                elif isinstance(param, mocAst.Assignment) :
                    if not isinstance(param.lvalue, mocAst.ID):
                        raise ValueError("Parameter must be a identifier")
                    # translate this parameter
                    name = param.lvalue.name
                    code_list = self.translate(param)
                    # mark that we have a default parameter
                    hasdefault = True
                else :
                    raise ValueError("Parameter must be a identifier")
                if name in node.parent.param_code :
                    raise ValueError("Already have a parameter named '%s'" % name)
                node.parent.param_code.append((name, code_list))
            return
        l = len(node.params)
        if isinstance(node.parent, mocAst.FuncCall) :
            # if this node is child of FuncCall, then we return
            # the parameters and count
            retlist = [ param.accept(self) for param in node.params ]
            return reversed(retlist), l
        if isinstance(node.parent, mocAst.Assignment) :
            temp = newtemp()
            if node.actionT == "list" :
                for param in node.params :
                    ret = param.accept(self)
                    self.emit(mocCst.Listitem(ret))
                self.emit(mocCst.Newlist(l, temp))
            else :
                for param in node.params :
                    param.actionT = "dict"
                    param.accept(self)
                self.emit(mocCst.Newdict(l, temp))
            return temp
        if isinstance(node.parent, mocAst.CollectionRef) :
            # when a collection reference, there must only one element in brackets
            if l > 1 : raise ValueError("Illegal expression in collection reference! % s" % node.coord)
            node.params[0].actionT = "list"
            return node.params[0].accept(self)
        # When a single code like "[1,2,3]" or "{1:xx,2:yy}" appears, we must know this is a kind of
        # create a list or dictionary, and may happen when pass a list to a callable
        if hasattr(node, "actionT"):
            temp = newtemp()
            mocAst.Assignment("=", mocAst.ID(temp), node).accept(self)
            return temp

    def visit_CollectionRef(self, node):
        # resolve the name of a collection object
        name = node.name.accept(self)
        temp = newtemp()
        if hasattr(node, "assign") and node.assign == True :
            subscript = node.subscript.accept(self)
            return self.emit(mocCst.Adddictkey(subscript, name))
        elif isinstance(node.subscript.params[0], mocAst.Assignment):
            node.subscript.accept(self)
            self.emit(mocCst.Slice(name, temp))
        else:
            subscript = node.subscript.accept(self)
            self.emit(mocCst.Collref(name, subscript))
            self.emit(mocCst.Getret(temp))
        return temp

    def visit_ObjectRef(self, node):
        # resolve the name of a object
        name = node.name.accept(self)
        field = node.field.accept(self)
        if hasattr(node, "assign") and node.assign == True :
            # when a object reference occurs at left hand side of a assignment
            # we may add some field dynamically
            self.emit(mocCst.Objassign(node.temp, name, field))
        else:
            temp = newtemp()
            self.emit(mocCst.Objref(name, field, temp))
            return temp

    def visit_FuncCall(self, node):
        # resolve the function name firstly
        func_name = node.name.accept(self)
        count = 0
        if node.args :
            # resolve the parameters
            params, count = node.args.accept(self)
            for param in params :
                self.emit(mocCst.Param(param))
        self.emit(mocCst.Call(func_name, count))
        # save the return value
        temp = newtemp()
        # use the "getret" instruction to get returned value
        self.emit(mocCst.Getret(temp))
        return temp

    def visit_If(self, node):
        node.cond.trueL = self.newlabel()
        node.cond.falseL = self.newlabel()
        ret = node.cond.accept(self)
        # if node.cond is a RelOp or a LogicalOp, then it would return nothing
        # if node.cond is a ID or a Constant, then we make a RelOp, to parse the If expression
        # ret is not a instance of mocTrue
        if ret is mocFalse :
            self.emit(mocCst.Goto(node.cond.falseL))
        elif ret and ret != mocTrue :
            # generate a line of compare statement
            cond = mocAst.RelOp("==", node.cond, mocAst.Constant("bool", True), node.cond.coord)
            cond.trueL = node.cond.trueL; cond.falseL = node.cond.falseL
            cond.accept(self)
            self.emitlabel(node.cond.trueL)
        else :
            self.emitlabel(node.cond.trueL)
        # translate the ifture branch
        node.iftrue.accept(self)
        if not node.iffalse :
            # for a If without Else statement
            self.emitlabel(node.cond.falseL)
        else :
            # has a Else statement
            next = self.newlabel()
            # add a label to jump to the end of Else
            self.emit(mocCst.Goto(next))
            self.emitlabel(node.cond.falseL)
            node.iffalse.accept(self)
            self.emitlabel(next)

    def visit_While(self, node):
        # generate two labels marking the limits of the While body
        begin = node.stmt.begin = self.newlabel()
        node.cond.trueL = self.newlabel()
        node.cond.falseL = node.stmt.end = self.newlabel()
        self.emitlabel(begin)
        if not isinstance(node.cond, (mocAst.RelOp, mocAst.LogicalOp)) :
            dummytest = mocAst.RelOp("==", node.cond, mocAst.Constant("bool", "true"), node.cond.coord)
            dummytest.trueL = node.cond.trueL
            dummytest.falseL = node.cond.falseL
            dummytest.accept(self)
        else :
            node.cond.accept(self)
        self.emitlabel(node.cond.trueL)
        node.stmt.accept(self)
        self.emit(mocCst.Goto(begin))
        self.emitlabel(node.cond.falseL)

    def visit_For(self, node):
        # convert the For iteration to While iteration, 
        # then translate the While statement
        if node.init :
            # we may encounter this situation for(: condition; next){}
            # so, if For has a initialization we translate it
            node.init.accept(self)
        # when the condition is empty, we generate a True constant
        cond = node.cond
        if not cond :
            cond = mocAst.Constant("bool", "true", node.coord)
        stmts = node.stmt.stmts[:]
        # adds the Next expression to the end of statements
        stmts.append(node.next)
        compound = mocAst.Compound(stmts, node.stmt.coord)
        compound.parent = node
        for2while = mocAst.While(cond, compound)
        # translate the While statement
        for2while.accept(self)

    def visit_ForEach(self, node):
        if not isinstance(node.ref, mocAst.ID) :
            raise ValueError("reference must be a variable in ForEach %s" % node.coord)
        # make a initial call for iterator
        obj = mocAst.ObjectRef(node.list, mocAst.ID("begin_itr"), node.coord)
        func0 = mocAst.FuncCall(obj, None, node.coord)
        # generate a reference call of list's has_next() method
        obj = mocAst.ObjectRef(node.list, mocAst.ID("has_next"), node.coord)
        func = mocAst.FuncCall(obj, None, node.coord)
        # generate a reference call of list's next() method
        obj = mocAst.ObjectRef(node.list, mocAst.ID("next"), node.coord)
        func2 = mocAst.FuncCall(obj, None, node.coord)
        obj.parent = func2.parent = node
        # use a assignment statement to assign a element to reference
        ref = mocAst.Assignment("=", node.ref, func2)
        stmts = [ref] + node.stmt.stmts[:]
        compound = mocAst.Compound(stmts, node.stmt.coord)
        compound.parent = node
        for2while = mocAst.While(func, compound)
        # translate the While statement
        # 1. perform initial call
        func0.accept(self)
        # 2. start iteration
        for2while.accept(self)
        # reconstruct the subtree
        node.list.parent = node

    def visit_Continue(self, node):
        p = node.parent
        count = 0
        while p and not hasattr(p, "begin"):
            p = p.parent
            count += 1
        if not p : raise ValueError("continue must in a loop or iteration % s" % node.coord)
        for i in xrange(count) :
            self.emit(mocCst.Leaveenv())
        self.emit(mocCst.Goto(p.begin))

    def visit_Break(self, node):
        p = node.parent
        count = 0
        while p and not hasattr(p, "end"):
            p = p.parent
            count += 1
        if not p : raise ValueError("break must in a loop or iteration % s" % node.coord)
        for i in xrange(count) :
            self.emit(mocCst.Leaveenv())
        self.emit(mocCst.Goto(p.end))

    def visit_Return(self, node):
        if node.expr :
            ret = mocAst.Assignment("=", mocAst.ID(newtemp()), node.expr).accept(self)
        else:
            ret = mocGNone
        p = node.parent
        count = -1
        while p and not hasattr(p, "ret_point"):
            p = p.parent
            count += 1
        if not p : raise ValueError("return must in a callable instance % s" % node.coord)
        for i in xrange(count) :
            self.emit(mocCst.Leaveenv())
        self.emit(mocCst.Setret(ret))
        self.emit(mocCst.Goto(p.ret_point))

    def visit_FuncDef(self, node):
        if node.name :
            funcdef = mocAst.FuncDef(None, node.params, node.body)
            closure = mocAst.Assignment("=", node.name, funcdef)
            return closure.accept(self)
        # the inherit attribute param_code is a dictionary
        # of name -> (name, code), for we may deal with parameter has default value
        node.param_code = []
        self.translate(node.params)
        # generate a label, so the return can jump to
        label = mocCst.Label(-1, 0)
        node.body.ret_point = label
        # translate the function body get the code list
        code_list = self.translate(node.body)
        label.lineno = len(code_list) - 1
        return node.param_code, code_list

    def visit_Compound(self, node):
        # create a new namespace
        self.emit(mocCst.Enterenv())
        self.translate_children(node)
        # leave the namespace
        self.emit(mocCst.Leaveenv())

    def visit_FileAST(self, node):
        self.emit(mocCst.Begin())
        self.translate_children(node)
        # mark the end of file
        self.emit(mocCst.End())
